import os
import torch
import torchvision
from torchvision import datasets, transforms
from torch.utils.data import Dataset
import numpy as np


transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
# 加载 MNIST 数据集
train_dataset = torchvision.datasets.MNIST(root='../data', train=True, download=True, transform=transform)
test_dataset = torchvision.datasets.MNIST(root='../data', train=False, download=True, transform=transform)

for i in range(1, 11):
    os.makedirs(f'./data/user{i}/train', exist_ok=True)
    os.makedirs(f'./data/user{i}/test', exist_ok=True)

def save_client_data(dataset, is_train=True, num_clients=10):
    data_per_client = len(dataset) // num_clients
    indices = torch.randperm(len(dataset))
    for i in range(num_clients):
        start_idx = i * data_per_client
        end_idx = start_idx + data_per_client
        client_indices = indices[start_idx:end_idx]
        images = [dataset[j][0].numpy() for j in client_indices]
        labels = [dataset[j][1] for j in client_indices]
        data_type = 'train' if is_train else 'test'
        np.savez(f'./data/user{i+1}/{data_type}/user{i+1}_{data_type}_mnist1.npz', images=images, labels=labels)

# 保存训练集和测试集
save_client_data(train_dataset, is_train=True, num_clients=10)
save_client_data(test_dataset, is_train=False, num_clients=10)
